/// 将STL文件转成DICOM文件
///
#include <unistd.h>

#include <itkVTKImageToImageFilter.h>

#include <vtkImageData.h>
#include <vtkSTLReader.h>
#include <vtkSmartPointer.h>
#include <vtkPointData.h>
#include <vtkPolyDataToImageStencil.h>
#include <vtkImageStencil.h>
#include <vtkMetaImageWriter.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>
#include <itkImageSeriesWriter.h>
#include <itkGDCMImageIO.h>
#include <itkNumericSeriesFileNames.h>



/// 将polydata转成 ImageData
vtkSmartPointer<vtkImageData> vtkPolyDataToImageData(vtkPolyData *polydata)
{
    vtkSmartPointer<vtkImageData> imageData = vtkSmartPointer<vtkImageData>::New();

    double bounds[6];
    polydata->GetBounds(bounds);
    bounds[0] -= 2.0;
    bounds[1] += 2.0;
    bounds[2] -= 2.0;
    bounds[3] += 2.0;
    bounds[4] -= 2.0;
    bounds[5] += 2.0;

    double spacing[3]; // desired volume spacing
    spacing[0] = 0.2;
    spacing[1] = 0.2;
    spacing[2] = 0.2;
    imageData->SetSpacing(spacing);

    // compute dimensions
    int dim[3];
    for (int i = 0; i < 3; i++) {
        dim[i] = static_cast<int>(ceil((bounds[i * 2 + 1] - bounds[i * 2]) / spacing[i]));
    }
    imageData->SetDimensions(dim);
    imageData->SetExtent(0, dim[0] - 1, 0, dim[1] - 1, 0, dim[2] - 1);

    double origin[3];
    origin[0] = bounds[0];
    origin[1] = bounds[2];
    origin[2] = bounds[4];
    // 这样和原始图偏移了 （-2，-2，-2）
    imageData->SetOrigin(origin);
    imageData->AllocateScalars(VTK_SHORT,1);

    // fill the image with foreground voxels:
    unsigned char inval = 255;
    unsigned char outval = 0;
    vtkIdType count = imageData->GetNumberOfPoints();
    for (vtkIdType i = 0; i < count; ++i) {
        imageData->GetPointData()->GetScalars()->SetTuple1(i, inval);
    }

    //polygonal data --> image stencil:
    vtkSmartPointer<vtkPolyDataToImageStencil> pdtoImageStencil = vtkSmartPointer<vtkPolyDataToImageStencil>::New();
    pdtoImageStencil->SetInputData(polydata);
    pdtoImageStencil->SetOutputOrigin(origin);
    pdtoImageStencil->SetOutputSpacing(spacing);
    pdtoImageStencil->SetOutputWholeExtent(imageData->GetExtent());
    pdtoImageStencil->Update();

    //cut the corresponding white image and set the background:
    vtkSmartPointer<vtkImageStencil> imageStencil = vtkSmartPointer<vtkImageStencil>::New();
    imageStencil->SetInputData(imageData);
    imageStencil->SetStencilConnection(pdtoImageStencil->GetOutputPort());
    imageStencil->ReverseStencilOff();
    imageStencil->SetBackgroundValue(outval);
    imageStencil->Update();

    imageData->DeepCopy(imageStencil->GetOutput());

    return imageData;
}

constexpr unsigned int Dimension = 3;
using PixelType = signed short;
using ImageType = itk::Image<PixelType, Dimension>;


bool writeToMhdFile(const char* stlPathName, const char *resultMhdFileName)
{
    if (0 != access(stlPathName, R_OK)) {
        return false;
    }

    vtkSmartPointer<vtkSTLReader> reader = vtkSmartPointer<vtkSTLReader>::New();
    reader->SetFileName(stlPathName);
    reader->Update();

    vtkSmartPointer<vtkPolyData> polydata = reader->GetOutput();
    vtkSmartPointer<vtkImageData> imageData = vtkPolyDataToImageData(polydata);

    vtkSmartPointer<vtkMetaImageWriter> writer =
      vtkSmartPointer<vtkMetaImageWriter>::New();
    writer->SetFileName(resultMhdFileName);
    writer->SetInputData(imageData);
    writer->Write();

    cout << "Success to convert " << stlPathName << " to " << resultMhdFileName << endl;
    return true;
}

bool convMhd2Dicom(ImageType::Pointer myitkImage, const char *dicomFolder)
{
    // GDCM
    using ImageIOType = itk::GDCMImageIO;
    using NamesGeneratorType = itk::NumericSeriesFileNames;
    ImageIOType::Pointer gdcmIO = ImageIOType::New();

    // write
    itksys::SystemTools::MakeDirectory(dicomFolder);


    using OutputPixelType = PixelType;
    constexpr unsigned int OutputDimension = 2;
    using Image2DType = itk::Image<OutputPixelType, OutputDimension>;

    using SeriesWriterType = itk::ImageSeriesWriter<ImageType, Image2DType>;

    NamesGeneratorType::Pointer namesGenerator = NamesGeneratorType::New();

    itk::MetaDataDictionary & dict = gdcmIO->GetMetaDataDictionary();
    std::string               tagkey, value;
    tagkey = "0008|0060"; // Modality
    value = "MR";
    itk::EncapsulateMetaData<std::string>(dict, tagkey, value);
    tagkey = "0008|0008"; // Image Type
    value = "DERIVED\\SECONDARY";
    itk::EncapsulateMetaData<std::string>(dict, tagkey, value);
    tagkey = "0008|0064"; // Conversion Type
    value = "DV";
    itk::EncapsulateMetaData<std::string>(dict, tagkey, value);


    SeriesWriterType::Pointer seriesWriter = SeriesWriterType::New();

    seriesWriter->SetInput(myitkImage);
    seriesWriter->SetImageIO(gdcmIO);


    ImageType::RegionType region =
            myitkImage->GetLargestPossibleRegion();

    ImageType::IndexType start = region.GetIndex();
    ImageType::SizeType  size = region.GetSize();


    std::string format = dicomFolder;

    format += "/image%03d.dcm";

    namesGenerator->SetSeriesFormat(format.c_str());

    namesGenerator->SetStartIndex(start[2]);
    namesGenerator->SetEndIndex(start[2] + size[2] - 1);
    namesGenerator->SetIncrementIndex(1);

    seriesWriter->SetFileNames(namesGenerator->GetFileNames());

    try
    {
        seriesWriter->Update();
    }
    catch (const itk::ExceptionObject & excp)
    {
        std::cerr << "Exception thrown while writing the series " << std::endl;
        std::cerr << excp << std::endl;
        return false;
    }

    cout << "Success to convert " << dicomFolder << endl;
    return true;
}

bool convMhd2Dicom(const char *mhdFileName, const char *dicomFolder)
{
    using ReaderType = itk::ImageFileReader<ImageType>;

    ReaderType::Pointer  reader = ReaderType::New();
    reader->SetFileName(mhdFileName);
    try
    {
        reader->Update();
    }
    catch (const itk::ExceptionObject & excp)
    {
        std::cerr << "Exception thrown while writing the image" << std::endl;
        std::cerr << excp << std::endl;
        return false;
    }

    return convMhd2Dicom(reader->GetOutput(), dicomFolder);
}

/// 将 stl 文件 转成 dicom文件
bool convStl2Dicom(const char *stlFileName, const char *dicomFolder)
{
    vtkSmartPointer<vtkSTLReader> reader = vtkSmartPointer<vtkSTLReader>::New();
    reader->SetFileName(stlFileName);
    reader->Update();

    vtkSmartPointer<vtkPolyData> polydata = reader->GetOutput();
    vtkSmartPointer<vtkImageData> imageData = vtkPolyDataToImageData(polydata);

    using FilterType = itk::VTKImageToImageFilter<ImageType>;
    auto filter = FilterType::New();
    filter->SetInput(imageData);

    try
    {
        filter->Update();
    }
    catch (const itk::ExceptionObject & error)
    {
        std::cerr << "Error: " << error << std::endl;
        return false;
    }

    return convMhd2Dicom(filter->GetOutput(), dicomFolder);
}

void showUsage(char *name)
{
    printf("Usage: %s option STL_file DICOM_folder/file.mhd\n", name);
    printf("       option: 1 - convert stl to mhd.\n");
    printf("               2 - convert mhd to dicom.\n");
    printf("               3 - convert stl to dicom.\n");
}

int main(int argc, char *argv[])
{
    if (argc != 4) {
        showUsage(argv[0]);
        return -1;
    }

    int option = atoi(argv[1]);
    char * inputFileName = argv[2];
    char * outputFileName = argv[3];

    switch (option) {
    case 1:
    {
        bool ret = writeToMhdFile(inputFileName, outputFileName);
        if (!ret) {
            showUsage(argv[0]);
            return -2;
        }
    }
        break;
    case 2:
    {
        convMhd2Dicom(inputFileName, outputFileName);
    }
        break;
    case 3:
    {
        convStl2Dicom(inputFileName, outputFileName);
    }
        break;
    default:
        break;
    }

    return 0;
}
